package ai.onnxruntime.example.objectdetection

import ai.onnxruntime.*
import ai.onnxruntime.extensions.OrtxPackage
import android.annotation.SuppressLint
import android.graphics.Bitmap
import android.graphics.BitmapFactory
import android.graphics.Canvas
import android.graphics.Color
import android.graphics.Paint
import android.graphics.PorterDuff
import android.graphics.PorterDuffXfermode
import android.os.Bundle
import android.util.Log
import android.widget.Button
import android.widget.ImageView
import android.widget.Toast
import androidx.appcompat.app.AppCompatActivity
import androidx.camera.core.*
import androidx.camera.lifecycle.ProcessCameraProvider
import androidx.camera.view.PreviewView
import android.content.pm.PackageManager
import android.Manifest
import android.graphics.Matrix
import android.media.ExifInterface
import androidx.core.app.ActivityCompat
import androidx.core.content.ContextCompat
import android.net.Uri
import java.io.ByteArrayInputStream
import java.io.ByteArrayOutputStream
import java.io.File
import java.io.InputStream
import android.os.Build
import android.content.ContentValues
import android.provider.MediaStore
import android.os.Environment
import android.content.Intent
import android.provider.Settings




class MainActivity : AppCompatActivity() {
    private var ortEnv: OrtEnvironment = OrtEnvironment.getEnvironment()
    private lateinit var ortSession: OrtSession
    private lateinit var inputImage: PreviewView
    private lateinit var outputImage: ImageView
    private lateinit var objectDetectionButton: Button
    private lateinit var detectionSaveButton: Button
    private var imageid = 0
    private lateinit var classes: List<String>
    private val CAMERA_PERMISSION_REQUEST_CODE = 1
    private lateinit var imageCapture: ImageCapture
    private lateinit var mutableBitmap: Bitmap

    companion object {
        const val TAG = "街景智能识别小助手"
    }

    @SuppressLint("UseCompatLoadingForDrawables")
    override fun onCreate(savedInstanceState: Bundle?) {
        super.onCreate(savedInstanceState)
        setContentView(R.layout.activity_main)

        inputImage = findViewById(R.id.previewView)
        outputImage = findViewById(R.id.imageView2)
        objectDetectionButton = findViewById(R.id.object_detection_button)
        detectionSaveButton = findViewById(R.id.detection_save_button)

        classes = readClasses()

        // Initialize Ort Session and register custom operators
        val sessionOptions: OrtSession.SessionOptions = OrtSession.SessionOptions()
        sessionOptions.registerCustomOpLibrary(OrtxPackage.getLibraryPath())
        ortSession = ortEnv.createSession(readModel(), sessionOptions)

        requestCameraPermission()

        objectDetectionButton.setOnClickListener {
            try {
                // Use the image captured by camera
                captureAndDetect()
                Toast.makeText(baseContext, "正在识别中···", Toast.LENGTH_SHORT)
                    .show()
            } catch (e: Exception) {
                Log.e(TAG, "Exception caught when performing ObjectDetection", e)
                Toast.makeText(baseContext, "没有找到可识别物体", Toast.LENGTH_SHORT).show()
            }
        }

        // 按钮点击事件监听器
        detectionSaveButton.setOnClickListener {
            try {
                if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.R) {
                    // 对于 Android 11（API 30）及以上
                    if (Environment.isExternalStorageManager()) {
                        saveImageToGallery(mutableBitmap)
                    } else {
                        requestStoragePermission()
                    }
                } else {
                    // 对于 Android 6.0 到 Android 10（API 23 到 29）
                    if (ContextCompat.checkSelfPermission(this, Manifest.permission.WRITE_EXTERNAL_STORAGE) == PackageManager.PERMISSION_GRANTED) {
                        saveImageToGallery(mutableBitmap)
                    } else {
                        requestStoragePermission()
                    }
                }
            } catch (e: Exception) {
                Toast.makeText(this, "保存失败", Toast.LENGTH_SHORT).show()
                Log.e(TAG, "Error saving image", e)
            }
        }

    }

    private fun requestCameraPermission() {
        if (ContextCompat.checkSelfPermission(this, Manifest.permission.CAMERA) != PackageManager.PERMISSION_GRANTED) {
            ActivityCompat.requestPermissions(
                this,
                arrayOf(Manifest.permission.CAMERA),
                CAMERA_PERMISSION_REQUEST_CODE
            )
        } else {
            startCamera()
        }
    }





    private fun startCamera() {
        val cameraProviderFuture = ProcessCameraProvider.getInstance(this)
        cameraProviderFuture.addListener(Runnable {
            val cameraProvider = cameraProviderFuture.get()

            val preview = Preview.Builder().build()
            val imageCapture = ImageCapture.Builder().build()

            val cameraSelector = CameraSelector.Builder()
                //.requireLensFacing(CameraSelector.LENS_FACING_BACK)
                .build()

            // Unbind all use cases before rebinding
            cameraProvider.unbindAll()

            // Bind preview and imageCapture to lifecycle
            cameraProvider.bindToLifecycle(this, cameraSelector, preview, imageCapture)

            // Set the surface provider for preview
            preview.setSurfaceProvider(inputImage.surfaceProvider)

            // Store the imageCapture instance so it can be used later
            this.imageCapture = imageCapture

            Log.d(TAG, "Camera is successfully bound and ready for capture.")
        }, ContextCompat.getMainExecutor(this))
    }
    
    // 捕获当前图像并进行目标检测
    private fun captureAndDetect() {
        if (!this::imageCapture.isInitialized) {
            Log.e(TAG, "ImageCapture is not initialized")
            return
        }

        val photoFile = File(getExternalFilesDir(null), "photo.jpg")
        val outputOptions = ImageCapture.OutputFileOptions.Builder(photoFile).build()

        imageCapture.takePicture(outputOptions, ContextCompat.getMainExecutor(this), object : ImageCapture.OnImageSavedCallback {
            override fun onImageSaved(outputFileResults: ImageCapture.OutputFileResults) {
                val savedUri = Uri.fromFile(photoFile)
                Log.d(TAG, "Photo saved successfully: $savedUri")
                performObjectDetectionFromCamera(savedUri)
            }

            override fun onError(exception: ImageCaptureException) {
                Log.e(TAG, "Photo capture failed: ${exception.message}", exception)
            }
        })
    }




    // Helper function to rotate the image based on Exif orientation
    private fun rotateImageIfNeeded(bitmap: Bitmap, orientation: Int): Bitmap {
        val matrix = Matrix()

        when (orientation) {
            ExifInterface.ORIENTATION_ROTATE_90 -> matrix.postRotate(90f)
            ExifInterface.ORIENTATION_ROTATE_180 -> matrix.postRotate(180f)
            ExifInterface.ORIENTATION_ROTATE_270 -> matrix.postRotate(270f)
            ExifInterface.ORIENTATION_NORMAL -> { /* No rotation needed */ }
            else -> { /* Handle other cases if necessary */ }
        }

        return Bitmap.createBitmap(bitmap, 0, 0, bitmap.width, bitmap.height, matrix, true)
    }


    private fun performObjectDetectionFromCamera(savedUri: Uri) {
        try {
            val inputStream = contentResolver.openInputStream(savedUri)
            if (inputStream != null) {
                var bitmap = BitmapFactory.decodeStream(inputStream)
                Log.d(TAG, "Bitmap loaded: ${bitmap.width}x${bitmap.height}")  // 打印 Bitmap 大小
                val scaledBitmap = Bitmap.createScaledBitmap(bitmap, 1280, 960, true)

                val exifInterface = ExifInterface(savedUri.path!!)
                val orientation = exifInterface.getAttributeInt(ExifInterface.TAG_ORIENTATION, ExifInterface.ORIENTATION_NORMAL)

                // Rotate the bitmap if needed
                bitmap = rotateImageIfNeeded(bitmap, orientation)

                outputImage.setImageBitmap(scaledBitmap)  // 显示捕获的图像

                // Convert Bitmap to InputStream
                val bitmapInputStream = bitmapToInputStream(bitmap)


                val objDetector = ObjectDetector()
                val result = objDetector.detect(bitmapInputStream, ortEnv, ortSession)

                updateUI(result)

                Toast.makeText(
                    baseContext,
                    "图像识别成功!",
                    Toast.LENGTH_SHORT
                ).show()
            }
            else{
                Log.e(TAG, "Failed to open input stream from URI: $savedUri")
            }
        } catch (e: Exception) {
            Log.e(TAG, "Error during object detection from camera", e)
            Toast.makeText(baseContext, "无法识别物品", Toast.LENGTH_SHORT).show()
        }
    }

    private fun updateUI(result: Result) {
        val mutableBitmap: Bitmap = result.outputBitmap.copy(Bitmap.Config.ARGB_8888, true)
        val canvas = Canvas(mutableBitmap)

        // 设置画笔用于绘制边框和文字
        val paint = Paint().apply {
            color = Color.WHITE// 固定边框和文字颜色为蓝色
            textSize = 108f
            xfermode = PorterDuffXfermode(PorterDuff.Mode.SRC_OVER)
        }

        val rectPaint = Paint().apply {
            color = Color.argb(255, 134, 48, 32)  // 固定边框颜色为红色
            style = Paint.Style.STROKE
            strokeWidth = 18f
        }

        val bgPaint = Paint().apply {
            color = Color.argb(200, 134, 48, 32)  // 半透明黑色用于文字背景
            style = Paint.Style.FILL
        }

        val boxit = result.outputBox.iterator()
        while (boxit.hasNext()) {
            val box_info = boxit.next()

            // 计算框的边界
            val left = box_info[0] - box_info[2] / 2
            val top = box_info[1] - box_info[3] / 2
            val right = box_info[0] + box_info[2] / 2
            val bottom = box_info[1] + box_info[3] / 2

            // 获取类别和置信度信息
            val category = classes[box_info[5].toInt()]
            val confidence = box_info[4]

            // 绘制矩形边框
            canvas.drawRect(left, top, right, bottom, rectPaint)

            // 绘制文字背景框
            val text = "%s: %.2f".format(category, confidence)
            val textWidth = paint.measureText(text)
            canvas.drawRect(left, top - 80, left + textWidth, top , bgPaint)

            // 绘制文字
            canvas.drawText(text, left, top - 20, paint)
        }

        // 将结果设置到 ImageView 中
        outputImage.setImageBitmap(mutableBitmap)
    }





    private val REQUEST_PERMISSION_WRITE_STORAGE = 1001

    private fun requestStoragePermission() {
        if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.R) {
            // 针对 Android 11（API 30）及以上版本
            if (Environment.isExternalStorageManager()) {
                // 已有管理所有文件的权限
                saveImageToGallery(mutableBitmap)
            } else {
                // 请求管理所有文件的权限
                try {
                    val intent = Intent(Settings.ACTION_MANAGE_APP_ALL_FILES_ACCESS_PERMISSION).apply {
                        data = Uri.parse("package:$packageName")
                    }
                    startActivity(intent)
                } catch (e: Exception) {
                    e.printStackTrace()
                    Toast.makeText(this, "无法打开权限设置页面", Toast.LENGTH_SHORT).show()
                }
            }
        } else {
            // 针对 Android 6.0 到 Android 10（API 23 到 29）
            if (ContextCompat.checkSelfPermission(this, Manifest.permission.WRITE_EXTERNAL_STORAGE) != PackageManager.PERMISSION_GRANTED) {
                ActivityCompat.requestPermissions(
                    this,
                    arrayOf(Manifest.permission.WRITE_EXTERNAL_STORAGE),
                    REQUEST_PERMISSION_WRITE_STORAGE
                )
            } else {
                // 已有存储权限
                saveImageToGallery(mutableBitmap)
            }
        }
    }


    override fun onRequestPermissionsResult(requestCode: Int, permissions: Array<String>, grantResults: IntArray) {
        super.onRequestPermissionsResult(requestCode, permissions, grantResults)

        when (requestCode) {
            CAMERA_PERMISSION_REQUEST_CODE -> {
                if (grantResults.isNotEmpty() && grantResults[0] == PackageManager.PERMISSION_GRANTED) {
                    startCamera()
                } else {
                    Toast.makeText(this, "需要摄像头权限", Toast.LENGTH_SHORT).show()
                }
            }

            REQUEST_PERMISSION_WRITE_STORAGE -> {
                if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.R) {
                    // 针对 Android 10+ 使用 MANAGE_EXTERNAL_STORAGE 权限
                    if (Environment.isExternalStorageManager()) {
                        saveImageToGallery(mutableBitmap)
                    } else {
                        Toast.makeText(this, "存储权限未授予，请授予权限", Toast.LENGTH_SHORT).show()
                    }
                } else {
                    if (grantResults.isNotEmpty() && grantResults[0] == PackageManager.PERMISSION_GRANTED) {
                        saveImageToGallery(mutableBitmap)
                    } else {
                        Toast.makeText(this, "需要存储的权限", Toast.LENGTH_SHORT).show()
                    }
                }
            }
        }
    }


    // 保存图片到图库
    private fun saveImageToGallery(bitmap: Bitmap) {
        val contentValues = ContentValues().apply {
            put(MediaStore.Images.Media.DISPLAY_NAME, "ImageName.jpg")
            put(MediaStore.Images.Media.MIME_TYPE, "image/jpeg")
            put(MediaStore.Images.Media.DATE_ADDED, System.currentTimeMillis() / 1000)
            put(MediaStore.Images.Media.DATE_TAKEN, System.currentTimeMillis())
        }

        val resolver = contentResolver
        val uri = resolver.insert(MediaStore.Images.Media.EXTERNAL_CONTENT_URI, contentValues)

        uri?.let { imageUri ->
            try {
                resolver.openOutputStream(imageUri)?.use { outputStream ->
                    if (bitmap.compress(Bitmap.CompressFormat.JPEG, 100, outputStream)) {
                        Toast.makeText(this, "图像保存到相册", Toast.LENGTH_SHORT).show()
                    } else {
                        Toast.makeText(this, "保存失败3", Toast.LENGTH_SHORT).show()
                    }
                }
            } catch (e: Exception) {
                e.printStackTrace()
                Toast.makeText(this, "保存失败2", Toast.LENGTH_SHORT).show()
            }
        } ?: run {
            Toast.makeText(this, "Failed to create new MediaStore record", Toast.LENGTH_SHORT).show()
        }
    }


    private fun readModel(): ByteArray {
        val modelID = R.raw.yolov8n_with_pre_post_processing
        return resources.openRawResource(modelID).readBytes()
    }

    private fun readClasses(): List<String> {
        return resources.openRawResource(R.raw.classes).bufferedReader().readLines()
    }

    override fun onDestroy() {
        super.onDestroy()
        ortEnv.close()
        ortSession.close()
    }
}

    // Convert Bitmap to InputStream
    private fun bitmapToInputStream(bitmap: Bitmap): InputStream {
        val byteArrayOutputStream = ByteArrayOutputStream()
        bitmap.compress(Bitmap.CompressFormat.PNG, 100, byteArrayOutputStream) // Compress as JPEG or PNG
        val byteArray = byteArrayOutputStream.toByteArray()
        return ByteArrayInputStream(byteArray)
    }

